# -*- coding: utf-8 -*-
import pickle
from collections import Counter

import cv2
import numpy as np
import tensorflow as tf
from scipy.misc import imresize

from facenet import facenet
import align.detect_face

import threading

cap = cv2.VideoCapture(0)
font = cv2.FONT_HERSHEY_SIMPLEX

embedding_file = "./models/embedding.pk"
modeldir = 'models/20190619-150742'

# knn的k值
k = 5
# 安全阈值
r = 0.7

a = "a".encode("utf-8")
d = "d".encode("utf-8")
w = "w".encode("utf-8")
s = "s".encode("utf-8")
q = "q".encode("utf-8")
h = "h".encode("utf-8")


def face_align():
    print('Creating networks and loading parameters')
    with tf.Graph().as_default():
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=True))
        pnet, rnet, onet = align.detect_face.create_mtcnn(sess, None)

        minsize = 20  # minimum size of face
        threshold = [0.6, 0.7, 0.7]  # three steps's threshold
        factor = 0.709  # scale factor

        def find_face(img):
            if img.ndim == 2:
                img = facenet.to_rgb(img)
            img = img[:, :, 0:3]

            bounding_boxes, _ = align.detect_face.detect_face(img, minsize, pnet, rnet, onet, threshold,
                                                              factor)
            nrof_faces = bounding_boxes.shape[0]
            if nrof_faces > 0:
                det = bounding_boxes[:, 0:4]
                det_arr = []
                img_size = np.asarray(img.shape)[0:2]
                for i in range(nrof_faces):
                    det_arr.append(np.squeeze(det[i]))
                positions = []
                faces = []
                for i, det in enumerate(det_arr):
                    det = np.squeeze(det)
                    bb = np.zeros(4, dtype=np.int32)
                    bb[0] = np.maximum(det[0] - 32 / 2, 0)
                    bb[1] = np.maximum(det[1] - 32 / 2, 0)
                    bb[2] = np.minimum(det[2] + 32 / 2, img_size[1])
                    bb[3] = np.minimum(det[3] + 32 / 2, img_size[0])
                    cropped = img[bb[1]:bb[3], bb[0]:bb[2], :]
                    scaled = imresize(cropped, (160, 160), interp='bilinear')
                    positions.append(bb)
                    faces.append(scaled)
                return positions, faces
            else:
                return None, None
        print("align model ready")
        return find_face


def video_face_recognition(embedding_path, modeldir):
    print("load face recognition model")
    with open(embedding_path, "rb") as fp:
        embedding_set = pickle.load(fp)
    image_size = 160

    with tf.Graph().as_default():
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=True))

        facenet.load_model(modeldir, sess)
        images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0")
        embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0")
        phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0")

        print('facenet embedding模型建立完毕')

        def face_recognition(img):
            img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_CUBIC)
            img = facenet.prewhiten(img)
            img = img.reshape(-1, image_size, image_size, 3)

            embedding_ = \
                sess.run(embeddings, feed_dict={images_placeholder: img, phase_train_placeholder: False})[0]

            img_names = []
            img_distence = []
            for name, emb in embedding_set:
                dist = np.sqrt(np.sum(np.square(embedding_ - emb)))
                img_distence.append(dist)
                img_names.append(name)

            # 为了加快速度，假设距离值互不相同
            sorted_list = sorted(img_distence)
            prop_names = []
            name_dis = dict()
            for d in sorted_list[0:k]:
                prop_names.append(img_names[img_distence.index(d)])
                if prop_names[-1] not in name_dis:
                    name_dis[prop_names[-1]] = []
                name_dis[prop_names[-1]].append(d)

            cnt = Counter(prop_names).most_common(1)
            name = cnt[0][0]
            dis = sorted(name_dis[name])[0]

            print(name, dis)
            return name, dis

        print("face recognition model ready")
        return face_recognition


def move_to_center(position, ser):
    if position and len(position) > 0:
        bb = position[0]
        x, y = (bb[2] - bb[0]) / 2 + bb[0], (bb[3] - bb[1]) / 2 + bb[1]
        print(x)
        if x > (320 + 10):
            print("move d")
            ser.write(a)
        elif x < (320 - 10):
            print("move a")
            ser.write(d)
        elif y > (240 + 10):
            print("move w")
            ser.write(w)
        elif y < (240 - 10):
            print("move s")
            ser.write(s)
        else:
            ser.write(q)
    else:
        ser.write(q)


def main():
    find_face = face_align()
    embedding = video_face_recognition(embedding_file, modeldir)
    import time
    fr = 0
    start = time.time()
    import serial
    ser = serial.Serial()
    ser.port = 'COM3'
    ser.baudrate = 115200
    ser.timeout = 0.2
    ser.open()
    while True:
        ret, frame = cap.read()
        image1 = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = cv2.resize(frame, (image1.shape[1] * 2, image1.shape[0] * 2))
        positions, faces = find_face(image1)
        names = []
        move_to_center(positions, ser)
        if positions and faces:
            for face in faces:
                name, dis = embedding(face)
                if dis < r:
                    names.append(name)
                else:
                    names.append("unknown")
            try:
                for name, bb in zip(names, positions):
                    cv2.rectangle(frame, (bb[0]*2, bb[1]*2), (bb[2]*2, bb[3]*2), (0, 255, 0), 5)
                    cv2.putText(frame, name, (bb[0]*2, bb[1]*2), font, 1, (200, 100, 255), 2, cv2.LINE_AA)
            except:
                pass
        cv2.imshow("face_detect", frame)
        fr += 1
        if time.time() - start > 5:
            print("帧率：", int(fr / 5))
            start = time.time()
            fr = 0

        if cv2.waitKey(1) & 0xFF == ord('q'):
            cap.release()
            cv2.destroyAllWindows()
            ser.close()
            break


if __name__ == '__main__':
    main()
